# -*- coding: utf-8 -*-
# ----------------------------
# @Time    : 2022/6/18 4:44 PM
# @Author  : changqingai
# @FileName: tf_sort_topk_accuracy.py
# ----------------------------


import tensorflow as tf

# ******* sort ***********
a = tf.random.shuffle(tf.range(6))
print("a:", a.numpy())
b = tf.sort(a, direction="DESCENDING")
print("b:", b.numpy())

idx = tf.argsort(a, direction="DESCENDING")
print("idx:", idx.numpy())
print("a:", a.numpy())

c = tf.gather(a, idx)
print("c:", c.numpy())

# ******* topk ***********
a = tf.random.shuffle(tf.range(6), seed=1)
print("a:", a.numpy())
b = tf.math.top_k(a, 3)
print("b:", b)
print("b.value:", b.values)
print("b.idx", b.indices)

# ******* topk 多维度 ***********
sample_num = 8
pred = tf.random.normal([sample_num, 10], stddev=1, mean=0)
true_label = tf.random.uniform([sample_num, ], minval=0, maxval=10, dtype=tf.int32)

pred = tf.math.top_k(pred, k=3)
print("pred.values:", pred.values)
print("pred.indices:", pred.indices)
pred = tf.transpose(pred.indices, [1, 0])
true_label = tf.broadcast_to(true_label, pred.shape)

correct = tf.equal(true_label, pred)
acc = tf.reduce_sum(tf.cast(correct, tf.int32)) / true_label.shape[1]
print("acc:", acc)


def topk_accuracy(pred, true_label, topk=[1,]):
    maxk = max(topk)
    sample_num = true_label.shape[0]

    pred = tf.math.top_k(pred, maxk)
    pred = tf.transpose(pred.indices, [1, 0])
    true_label = tf.broadcast_to(true_label, pred.shape)

    correct = tf.equal(pred, true_label)
    res = []
    for k in topk:
        correct_k = tf.cast(correct[:k], dtype=tf.float32)
        correct_k = tf.reduce_sum(correct_k)
        acc = float(correct_k) / sample_num
        res.append(acc)
    return res


sample_num = 8
pred = tf.random.normal([sample_num, 10], stddev=1, mean=0)
true_label = tf.random.uniform([sample_num, ], minval=0, maxval=10, dtype=tf.int32)

res = topk_accuracy(pred, true_label, topk=[1, 3, 5])
print("top_k.acc:", res)
